Persistent MXFP8 for experts layers in MoE - #1940
Conversation
…ate moe architecture spec to define emitted tensors
|
qq: how does the reward track vs the bf16 baseline here? also curious about logprobs mismatch |
There was a problem hiding this comment.
Code Review
This pull request introduces support for expert-only MXFP8 quantization on Blackwell GPUs using Transformer Engine and vLLM, including benchmark scripts, documentation, and updated sequence packing alignment logic. The review feedback highlights two key improvements: first, a potential shape mismatch crash in _load_batched_moe_fp8_tensor should be resolved by always unbinding and loading experts individually; second, the Blackwell hardware check in expert_mxfp8.py should be made more robust by verifying the major compute capability version is 10 rather than checking for specific minor versions.
| if param.shape[0] == loaded_weight.shape[0]: | ||
| success = weight_loader( | ||
| param, | ||
| loaded_weight, | ||
| target_name, | ||
| shard_id=shard_id, | ||
| expert_id=0, | ||
| return_success=True, | ||
| ) | ||
| if not success: | ||
| raise ValueError(f"Fused loading failed for batched MoE tensor {wire_name!r}") | ||
| return True | ||
|
|
||
| # Expert-parallel vLLM keeps only a subset locally. Retain the compact wire | ||
| # format, but let the loader map global expert IDs one view at a time. | ||
| loaded_any = False | ||
| for expert_id, expert_weight in enumerate(loaded_weight.unbind(0)): | ||
| loaded_any = ( | ||
| bool( | ||
| weight_loader( | ||
| param, | ||
| expert_weight, | ||
| target_name, | ||
| shard_id=shard_id, | ||
| expert_id=expert_id, | ||
| return_success=True, | ||
| ) | ||
| ) | ||
| or loaded_any | ||
| ) | ||
| if not loaded_any: | ||
| raise ValueError(f"No local expert accepted batched MoE tensor {wire_name!r}") | ||
| return True |
There was a problem hiding this comment.
In _load_batched_moe_fp8_tensor, when param.shape[0] == loaded_weight.shape[0], calling weight_loader with the 3D loaded_weight and expert_id=0 will cause a shape mismatch and crash in vLLM. vLLM's weight_loader for MoE expects a 2D tensor for a single expert and copies it to param[expert_id]. We should remove this branch and always use the loop to unbind and load each expert individually. This is robust and works correctly in all environments (including single-GPU or non-expert-parallel setups).
# Retain the compact wire format, but let the loader map global expert IDs one view at a time.
loaded_any = False
for expert_id, expert_weight in enumerate(loaded_weight.unbind(0)):
loaded_any = (
bool(
weight_loader(
param,
expert_weight,
target_name,
shard_id=shard_id,
expert_id=expert_id,
return_success=True,
)
)
or loaded_any
)
if not loaded_any:
raise ValueError(f"No local expert accepted batched MoE tensor {wire_name!r}")
return True| capability = torch.cuda.get_device_capability() | ||
| if capability not in ((10, 0), (10, 3)): | ||
| raise RuntimeError(f"Expert MXFP8 requires SM100 or SM103, got SM{capability[0]}{capability[1]}") |
There was a problem hiding this comment.
Checking for specific Blackwell compute capabilities (10, 0) and (10, 3) is fragile and might break on other Blackwell GPU variants (e.g., GB200 or future revisions with minor version changes like 10.1 or 10.2). It is safer and more future-proof to check if the major version of the compute capability is 10.
| capability = torch.cuda.get_device_capability() | |
| if capability not in ((10, 0), (10, 3)): | |
| raise RuntimeError(f"Expert MXFP8 requires SM100 or SM103, got SM{capability[0]}{capability[1]}") | |
| capability = torch.cuda.get_device_capability() | |
| if capability[0] != 10: | |
| raise RuntimeError(f"Expert MXFP8 requires Blackwell (SM10x), got SM{capability[0]}{capability[1]}") |
|
Hi @erictang000 , I've pushed some changes which solve the perf degradations I was seeing (batched casting + sending, vllm settings) -- now the mxfp8 tps is significantly higher than the bf16 baseline (10%+ improvement ). TPS comparison: Comparisons for reward/logprob diff consistent with above. I've taken the PR out of draft, currently working on a bit of infra code to merge all quantization strategies (FP8, mxfp8, potentially nvfp4 later) under a unified spec + strategy class definition (similar idea to what you've mentioned here #1898 (comment)), will create a new PR for that + nvfp4 when complete. |




Building off of #1937 and #1898, introduce persistent MXFP8 quantization for expert layers in MoE
Reuses blockwise casting logic from #1898 to write E8M0 scale factors on 1x32 blocks & deferred load of FP8 weights so optimizer states initialized on FP32 master weights, and pack E8M0 scale factors into uint8 bytes for VLLM ModelOpt-MXFP8 offline quantization.
_iter_sync_tensorswhich callsiter_serialized_fp8_tensorsnow handles both FP8 (hopper-style) and MXFP8 (blackwell-style).Todo: I still observe perf degradations from mxfp8 vs the bf16 baseline, need to debug why
